{-# LANGUAGE RankNTypes, FlexibleInstances #-}

-- | Generalized versions of exception-related IO functions to work with any
-- monad transformer stack on top of IO.  Like MonadIO, but allows the IO
-- operation to unlift back into the monad.
--
-- Author: Stephen Blackheath
-- Copyright iPwn Studios Inc.
-- BSD License
module IPwn.Helper.Transformer where

import qualified Control.Concurrent as C
import qualified Control.Exception as E
import Control.Monad.State
import Data.IORef


class MonadFinally m where
    finally :: m a -> m b -> m a

class (MonadIO m, MonadFinally m) => MonadReliftIO m where
    -- | Lift an IO into this monad, passing it an 'unlift' function that it can
    -- use to get back to this monad.
    reliftIO :: ((forall a . m a -> IO a) -> IO b) -> m b
    -- | Generalized 'Control.Concurrent.forkIO' for any MonadReliftIO.
    forkIO :: m () -> m C.ThreadId
    forkIO code = reliftIO $ \unlift -> C.forkIO $ unlift code

instance MonadFinally IO where
    finally = E.finally

instance MonadReliftIO IO where
    reliftIO code = code id

instance MonadReliftIO m => MonadFinally (StateT s m) where
    finally = reliftIOFinally

instance MonadReliftIO m => MonadReliftIO (StateT s m) where
    -- Note! This assumes the function calling it is running a single thread
    -- only - true of catch and finally but not forkIO.
    reliftIO code = do
        st <- get
        ref <- liftIO $ newIORef st
        ret <- lift $ reliftIO $ \unlift ->
            code $ \op -> do  -- op :: StateT s m a -> IO a
                st <- readIORef ref
                (ret, st') <- unlift $ runStateT op st
                writeIORef ref st'
                return ret
        st' <- liftIO $ readIORef ref
        put st'
        return ret
    -- I don't think we're using this, but it should be correct anyway.
    -- reliftIO assumes single-threaded operation, so we need to write a
    -- stateT-specific version of forkIO.
    forkIO code = do
        st <- get
        lift $ forkIO $ evalStateT code st


-- | Generalized 'Control.Concurrent.catch' for any MonadReliftIO
catch :: (MonadReliftIO m, E.Exception e) => m a -> (e -> m a) -> m a
catch code handler = reliftIO $ \unlift ->
    unlift code `E.catch` \exc -> unlift $ handler exc

-- | Generalized 'Control.Exception.finally' for any MonadReliftIO. 
reliftIOFinally :: MonadReliftIO m =>
                   m a         -- ^ computation to run first
                -> m b         -- ^ computation to run afterward (even if an exception
                               -- was raised)
                -> m a         -- returns the value from the first computation
a `reliftIOFinally` sequel = reliftIO $ \unlift ->
    unlift a `E.finally` unlift sequel

